'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 3_study_support_function_call_with_prompt.py
* @Time: 2025/10/22
* @All Rights Reserve By Brtc
'''
import json
import os
from typing import Type, Any, Dict, TypedDict, Optional

import dotenv
import requests
from langchain_community.tools import GoogleSerperRun
from langchain_community.utilities import GoogleSerperAPIWrapper
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableConfig, RunnablePassthrough
from langchain_core.tools import BaseTool, render_text_description_and_args
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field

dotenv.load_dotenv()

class GaodeWeatherSchema(BaseModel):
    city:str = Field(description="需要查询的天气的城市，如：长沙")

class GoogleSerperSchema(BaseModel):
    query:str = Field(description="执行谷歌搜索的查询语句")

class GaodeWeatherTool(BaseTool):
    """根据传入的城市名称运行调用API 获取传入城市的天气预报信息"""
    name:str = "gaode_weather_tool"
    description:str = "当你想查询天气的时候可以使用这个工具."
    args_schema:Type[BaseModel] = GaodeWeatherSchema

    def _run(self, *args:Any, **kwargs:Any) -> Any:
        try:
            gaode_api_key = os.getenv("GAODE_API_KEY")
            gaode_api_url = os.getenv("GAODE_BASE_URL")
            if not  gaode_api_key or not gaode_api_url:
                return f"请配置正确的高德API KEY 和 URL"
            else:
                #1、从参数中获取城市
                city = kwargs.get("city")
                #2、行政区域的查询
                session = requests.session()
                city_response = session.request(
                    method="GET",
                    url = f"{gaode_api_url}/config/district?key={gaode_api_key}&keywords={city}&subdistrict=0",
                    headers= {"Content-Type": "application/json; charset=utf-8"}
                )
                city_response.raise_for_status()
                city_data = city_response.json()
                print(city_data)
                if city_data.get("info") == "OK":
                    ad_code = city_data.get("districts")[0]["adcode"]
                    weather_info = session.request(
                        method="GET",
                        url=f"{gaode_api_url}/weather/weatherInfo?key={gaode_api_key}&city={ad_code}&extensions=all",
                    )
                    weather_info.raise_for_status()
                    weather_data = weather_info.json()
                    if weather_data.get("info") == "OK":
                        return json.dumps(weather_data)
                return f"获取{city}天气预报失败！"
        except Exception as e:
            return f"获取天气失败"

#1、定义工具列表
gaode_weather = GaodeWeatherTool()
google_serper = GoogleSerperRun(
    name = "google_serper",
    description=("一个低成本的搜索工具！"),
    api_wrapper=GoogleSerperAPIWrapper(),
)

tool_dict = {
    gaode_weather.name:gaode_weather,
    google_serper.name:google_serper,
}

tools = [tool for tool in tool_dict.values()]


class ToolCallRequest(TypedDict):
    """工具调用请求字典"""
    name:str
    arguments:Dict=[str, Any]

def invoke_tool(
        tool_call_request:ToolCallRequest, config:Optional[RunnableConfig]=None
)->Any:
    """
    我们可以使用的执行工具调用函数
    :param tool_call_request:一个包含键名和参数的字典，名称必须与现有的工具名称匹配，参数是该工具的调用参数
    :param config:这是langchain使用的包含回调、元素据、等信息的配置信息
    :return 请求工具的输出内容
    """
    tool_name_to_tool = {tool.name:tool for tool in tools}
    name = tool_call_request["name"]
    requested_tool = tool_name_to_tool[name]
    return requested_tool.invoke(tool_call_request["arguments"], config=config)

#通过提示词去构建函数调用信息
system_prompt = """  
您是一个由OpenAI开发的聊天机器人，可以访问以下一组工具。
以下是每个工具的名称和描述：
{rendered_tools}
根据用户输入，返回要使用的工具的名称和输入。
将您的响应作为具有`name`和`arguments`键的JSON块返回。
`arguments`应该是一个字典，其中键对应于参数名称，值对应与请求的值。
"""
prompt = ChatPromptTemplate.from_messages([
    ("system", system_prompt),
    ("human", "{query}")
]).partial(rendered_tools= render_text_description_and_args(tools))#通过render_text_description_and_args将tools 的 信息转换成提示词细信息

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

#调用链
chain =(
    {"query":RunnablePassthrough()} # 提问占位符
    |prompt
    |llm
    |JsonOutputParser()
    |RunnablePassthrough.assign(output= invoke_tool)
)

print(chain.invoke("请写一个关于程序员的笑话？"))



